/*
 * MIT License
 *
 * Copyright (c) 2020 wen.gu <454727014@qq.com>
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 * SOFTWARE.
 */

 /***************************************************************************
 * Name: tiny_decoder.c
 *
 * Purpose: implementation a decoder with tiny decompress
 *
 * Developer:
 *   wen.gu , 2022-03-04
 *
 * TODO:
 *
 ***************************************************************************/

 /******************************************************************************
 **    INCLUDES
 ******************************************************************************/

#include "tiny_decoder.h"

#include <stdio.h>
#include <string.h>
#include "lz4/lz4.h"
/******************************************************************************
 **    MACROS
 ******************************************************************************/
#define LOGE printf
#define LOGD printf

#define Y_BLOCK_WIDTH 16
#define UV_BLOCK_WIDTH 8

#define Y_BLOCK_SIZE (Y_BLOCK_WIDTH * Y_BLOCK_WIDTH)
#define UV_BLOCK_SIZE (UV_BLOCK_WIDTH * UV_BLOCK_WIDTH)

#define RESIDUAL_STEP 2

/******************************************************************************
 **    VARIABLE DEFINITIONS
 ******************************************************************************/

/******************************************************************************
 **    inner FUNCTION DEFINITIONS
 ******************************************************************************/

static uint32_t get_coordinate_size(uint32_t width, uint32_t height, uint32_t& y_coordinate_size, uint32_t& uv_coordinate_size)
{
    uint32_t y_size = width * height;
    y_coordinate_size = y_size / Y_BLOCK_SIZE * 4; /** use 2 byte indicat coordinate, (x,y) total 4 byte */
    uv_coordinate_size = y_size / 2 / UV_BLOCK_SIZE * 4; /** use 2 byte indicat coordinate, (x,y) total 4 byte */
    return y_coordinate_size + uv_coordinate_size;
}


static void partition_decode(const uint8_t* ref_frame, const uint8_t* residual_frame, const uint16_t* coordinate_buf, uint8_t* out_frame, uint32_t width, uint32_t height, uint32_t block_width)
{
    for (uint32_t i = 0; i < height; i += block_width)
    {
        for (uint32_t j = 0; j < width; j += block_width)
        {
            uint32_t pos_x = *coordinate_buf++;
            uint32_t pos_y = *coordinate_buf++;
            for (size_t block_h = 0; block_h < block_width; block_h++)
            {
                for (size_t block_w = 0; block_w < block_width; block_w++)
                {
                    uint32_t ref_idx = (pos_y + block_h) * width + pos_x + block_w;
                    uint32_t out_idx = (i + block_h) * width + j + block_w;
                    out_frame[out_idx] = ref_frame[ref_idx] + residual_frame[out_idx] * RESIDUAL_STEP;
                }                
            }
        }
    }    
}

/******************************************************************************
 **    FUNCTION DEFINITIONS
 ******************************************************************************/
TinyDecoder::TinyDecoder()
{

}

TinyDecoder::~TinyDecoder()
{
    if (reference_frame_buf_)
    {
        delete[] reference_frame_buf_;
    }

    if (out_raw_buf_)
    {
        delete[] out_raw_buf_;
    }

}

bool TinyDecoder::setConfigure(KfaDecoderParam& param)
{    
    if (param.raw_type != kfa_raw_type::NV12)
    {
        LOGE("Tiny decoder only support the raw type of NV12\n");
        return false;
    }

    img_width_ = param.pic_width; //todo refine me
    img_height_ = param.pic_height;
    
    coordinate_buf_size_ = get_coordinate_size(img_width_, img_height_, y_coordinate_buf_size_, uv_coordinate_buf_size_);
    LOGD("picture size(w:%d, h:%d), coordinate size(total: %d, y:%d, uv:%d)\n", 
        img_width_, img_height_, coordinate_buf_size_, y_coordinate_buf_size_, uv_coordinate_buf_size_);
    return true;
}

bool TinyDecoder::decode(const uint8_t* in_data, uint32_t in_size, KfaBuffer& out_buf)
{
    //printf("[%s.%d]===>:\n", __FUNCTION__, __LINE__); 
    if (!out_raw_buf_)
    {/** inidicate to decode first frame, then just decompress */
        out_frame_size_ = img_width_ * img_height_ * 3 / 2; //only for NV12
        out_raw_size_ =  out_frame_size_ + coordinate_buf_size_;
        out_raw_buf_ = new uint8_t[out_raw_size_];
        predictive_buf_ = new uint8_t[out_frame_size_];

        //reference_frame_buf_ = new uint8_t[out_frame_size_];
        for (uint32_t i = 0; i < TINY_DEC_OUT_BUF_CNT; i++)
        {
            out_frame_buf_[i] = new uint8_t[out_frame_size_];
        }
        
        //out_buf.buf = new uint8_t[out_frame_size_];
        reference_frame_buf_ = out_frame_buf_[out_idx_++ % TINY_DEC_OUT_BUF_CNT];

        out_buf.buf = reference_frame_buf_;
        out_buf.release = [](KfaBuffer& buf)
        {
            //do nothing
        };

        int res_size = LZ4_decompress_safe((const char*)in_data,  (char*)out_buf.buf, in_size, out_frame_size_);

        if (res_size <= 0)
        {
            LOGE("decompress data failed(%d)\n", res_size);
            return false;
        }

        out_buf.size = res_size;
        return true;
    }
    //printf("[%s.%d]===>:\n", __FUNCTION__, __LINE__); 
    int res_size = LZ4_decompress_safe((const char*)in_data,  (char*)out_raw_buf_, in_size, out_raw_size_);

    if (res_size <= 0)
    {
        return false;
    }
    //printf("[%s.%d]===>:res_size: %d\n", __FUNCTION__, __LINE__, res_size); 
    uint16_t* y_coordinate_buf = (uint16_t*)out_raw_buf_;
    uint16_t* uv_coordinate_buf = (uint16_t*)(out_raw_buf_ + y_coordinate_buf_size_);
    uint8_t* residual_buf = out_raw_buf_ + coordinate_buf_size_;

    //printf("[%s.%d]===>:\n", __FUNCTION__, __LINE__); 
    /** do motion compensation */
    uint8_t* frame_buf = out_frame_buf_[out_idx_++ % TINY_DEC_OUT_BUF_CNT];
    uint32_t y_size = img_width_ * img_height_;
    partition_decode(reference_frame_buf_, residual_buf, y_coordinate_buf, frame_buf, img_width_, img_height_, Y_BLOCK_WIDTH);
    partition_decode(reference_frame_buf_ + y_size, residual_buf + y_size, uv_coordinate_buf, frame_buf + y_size, img_width_ / 2, img_height_, UV_BLOCK_WIDTH);

    reference_frame_buf_ = frame_buf;

    //printf("[%s.%d]===>:\n", __FUNCTION__, __LINE__); 
    out_buf.buf = frame_buf;
    out_buf.release = [](KfaBuffer& buf)
    {
        //todo something
    };
    //printf("[%s.%d]===>:\n", __FUNCTION__, __LINE__); 
    out_buf.size = out_frame_size_;        

    //LOGD("[%s.%d]===>:\n", __FUNCTION__, __LINE__);
    return true;
}

